import numpy as np


def padding_2d(array: np.ndarray, ks: tuple, mode: str, value: int or np.ndarray = 0):
    """Extend the 2D `array` using the specified boundary `mode`, w.r.t. kernel size `ks`.

    `array`: a numpy ndarray in the shape of (h, w, ...).

    `ks`: size (kh, kw) of the kernel used in filtering.

    `mode`: the boundary mode specify how the array will be padded.

    `value`: a constant value that fills the padding in constant mode.
    """

    assert mode in ["constant", "replicate", "symmetric", "circular"]

    ah, aw = array.shape[:2]
    kh, kw = ks
    ph = ah + 2 * (kh - 1)
    pw = aw + 2 * (kw - 1)

    im = np.zeros(shape=(ph, pw) + array.shape[2:], dtype=array.dtype)
    im[kh - 1 : ah + kh - 1, kw - 1 : aw + kw - 1] = array

    if mode == "constant":
        im[: kh - 1, :] = value
        im[ah + kh - 1 :, :] = value
        im[:, : kw - 1] = value
        im[:, aw + kw - 1 :] = value

    if mode == "replicate":
        im[: kh - 1, :] = np.expand_dims(im[kh - 1, :], axis=0)
        im[ah + kh - 1 :, :] = np.expand_dims(im[ah + kh - 2, :], axis=0)
        im[:, : kw - 1] = np.expand_dims(im[:, kw - 1], axis=1)
        im[:, aw + kw - 1 :] = np.expand_dims(im[:, aw + kw - 2], axis=1)

    if mode == "symmetric":
        im[: kh - 1, :] = im[2 * kh - 3 : kh - 2 : -1, :]
        im[ah + kh - 1 :, :] = im[ah + kh - 2 : ah - 1 : -1, :]
        im[:, : kw - 1] = im[:, 2 * kw - 3 : kw - 2 : -1]
        im[:, aw + kw - 1 :] = im[:, aw + kw - 2 : aw - 1 : -1]

    if mode == "circular":
        im[: kh - 1, :] = im[ah : ah + kh - 1, :]
        im[ah + kh - 1 :, :] = im[kh - 1 : 2 * kh - 2, :]
        im[:, : kw - 1] = im[:, aw : aw + kw - 1]
        im[:, aw + kw - 1 :] = im[:, kw - 1 : 2 * kw - 2]

    return im


def row_shape_windows(array: np.ndarray, ks: tuple):
    """Convert cells in a sliding window of size `ks` into a row, for each sliding
    window over the 2D `array`. Assume that the shape of `array` is (ah, aw, ...) and
    the kernel size `ks` is (kh, kw), the returned row-shape windows will be of the
    shape ((ah - kh + 1) * (aw - kw + 1), kh * kw).

    `array`: a numpy ndarray in the shape of (h, w, ...).

    `ks`: size (kh, kw) of the kernel used in filtering.
    """

    ah, aw = array.shape[:2]
    kh, kw = ks

    windows_shape = ((ah - kh + 1) * (aw - kw + 1), kh * kw)

    # Generate all possible row and column indices for sliding windows
    kh_indices, kw_indices = np.indices(dimensions=(kh, kw))

    ah_ind = np.arange(ah - kh + 1)[:, np.newaxis, np.newaxis, np.newaxis]
    kh_ind = kh_indices[np.newaxis, np.newaxis, :, :]
    aw_ind = np.arange(aw - kw + 1)[np.newaxis, :, np.newaxis, np.newaxis]
    kw_ind = kw_indices[np.newaxis, np.newaxis, :, :]

    windows = array[ah_ind + kh_ind, aw_ind + kw_ind]
    windows = np.reshape(windows, newshape=windows_shape + array.shape[2:])

    return windows


def spatial_filtering(array: np.ndarray, ks: tuple, kernel: "callable" = None):
    """Perform spatial filtering on the padded array, using the provided kernel.

    `array`: a numpy ndarray in the shape of (ah, aw, ...), which usually results
    from padding the original array using function `padding_2d`.

    `ks`: size (kh, kw) of the kernel used in filtering.

    `kernel`: a callable that receives an array of cells in a sliding window,
    performs filtering, and returns the result value for that cell.
    """

    ah, aw = array.shape[:2]
    kh, kw = ks

    row_windows = row_shape_windows(array, ks)
    row_pixels = np.array(list(map(kernel, row_windows)))
    target_shape = (ah - kh + 1, aw - kw + 1) + row_pixels.shape[1:]
    filtered_array = row_pixels.reshape(*target_shape)

    return filtered_array


__all__ = ["padding_2d", "row_shape_windows", "spatial_filtering"]
